"""
上传数据到Dify知识库模块
包含详细的日志记录和错误处理
基于 dify_utils.py 实现
"""
import re
import traceback
from pathlib import Path
from conn_manager import FAQ_DATA_DIR, EC_DATA_DIR, DIFY_KNOWLEDGE_DATA
from dify_utils import FileToDify
import logging

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def extract_product_name(filename):
    """
    从文件名提取产品名称
    支持三种格式：prodname-faq.txt, FAQ-prodname.txt, TroubleShooting-prodname.txt
    :param filename: 文件名
    :return: 产品名称，如果无法提取则返回None
    """
    filename = Path(filename).stem  # 去掉扩展名
    
    # 匹配三种模式
    patterns = [
        r'^(.+)-faq$',  # prodname-faq.txt
        r'^FAQ-(.+)$',  # FAQ-prodname.txt
        r'^TroubleShooting-(.+)$'  # TroubleShooting-prodname.txt
    ]
    
    for pattern in patterns:
        match = re.match(pattern, filename, re.IGNORECASE)
        if match:
            product_name = match.group(1)
            # logger.info(f"📋 从文件名 {filename} 提取产品名称: {product_name}")
            return product_name
    
    logger.warning(f"⚠️ 无法从文件名 {filename} 提取产品名称")
    return None

def check_document_exists(dify_client, knowledge_id, document_name):
    """
    检查文档是否存在的辅助函数
    """
    try:
        # 确保我们有正确的当前知识库信息
        if not dify_client.current_knowledge or dify_client.current_knowledge.get('id') != knowledge_id:
            # 如果当前知识库信息不匹配，重新查询知识库列表
            dify_client.query_knowledge_list()
            # 找到对应的知识库
            for name, data in dify_client.knowledge_list.items():
                if data.get('id') == knowledge_id:
                    dify_client.current_knowledge = data
                    break
        
        # 查询文档列表
        dify_client.query_document_list(knowledge_id)
        knowledge_name_for_search = dify_client.current_knowledge.get('name')
        
        if knowledge_name_for_search and knowledge_name_for_search in dify_client.queried_document:
            for doc_id, doc_data in dify_client.queried_document[knowledge_name_for_search].items():
                if doc_data.get('name') == document_name:
                    # logger.info(f"📄 找到现有文档: {document_name} (ID: {doc_id})")
                    return doc_id
        
        logger.info(f"📄 文档不存在，将创建新文档: {document_name}")
        return None
    except Exception as e:
        logger.warning(f"⚠️ 检查文档存在性时出错: {str(e)}")
        # 如果检查失败，假设文档不存在，尝试上传新文档
        return None

def upload_to_errorcode_dify(data_dir=EC_DATA_DIR, environments=['local', 'aws'], export_all=True, files_to_upload=None):
    """
    上传错误代码数据到Dify的errorcode知识库
    :param data_dir: 错误代码数据目录
    :param environments: 要上传的环境列表
    :param export_all: 是否全量上传，False时进行增量上传
    :param files_to_upload: 增量上传时的文件列表
    :return: 上传结果统计
    """
    if export_all:
        logger.info("🔧 开始全量上传错误代码数据到Dify...")
    else:
        logger.info("🎯 开始增量上传错误代码数据到Dify...")
    
    results = {'successful': 0, 'total': 0, 'errors': []}
    
    # 根据export_all参数决定上传文件列表
    if export_all:
        # 全量上传：扫描目录获取所有文件
        data_path = Path(data_dir)
        if not data_path.exists():
            logger.error(f"❌ 错误代码数据目录不存在: {data_path}")
            return results
            
        txt_files = list(data_path.glob("*.txt"))
        if not txt_files:
            logger.warning(f"⚠️ 错误代码数据目录中没有找到txt文件: {data_path}")
            return results
            
        logger.info(f"📋 找到 {len(txt_files)} 个错误代码文件（全量上传）")
        # 转换为文件信息格式
        file_list = [{'file_name': f.name, 'file_path': str(f), 'file_size': f.stat().st_size} for f in txt_files]
    else:
        # 增量上传：使用提供的文件列表
        if not files_to_upload:
            logger.info("😌 没有错误代码文件需要增量上传")
            return results
            
        # 过滤出错误代码文件
        errorcode_files = [f for f in files_to_upload if 'errorcode' in f['file_path'].lower()]
        if not errorcode_files:
            logger.info("😌 没有错误代码文件需要增量上传")
            return results
            
        logger.info(f"🎯 找到 {len(errorcode_files)} 个错误代码文件需要增量上传")
        file_list = errorcode_files
    
    # 遍历每个环境
    for env in environments:
        try:
            logger.info(f"🌍 开始处理 {env} 环境...")
            
            # 创建FileToDify实例
            dify_client = FileToDify(url_data=DIFY_KNOWLEDGE_DATA[env])
            
            # 查询知识库列表
            dify_client.query_knowledge_list()
            
            # 创建或获取errorcode知识库 - 避免使用类方法，使用独立客户端
            knowledge_name = "errorcode"
            if knowledge_name not in dify_client.knowledge_list:
                logger.info(f"📚 创建新的错误代码知识库: {knowledge_name} ({env} 环境)")
                # 使用实例方法创建知识库，避免类变量污染
                success = dify_client.create_knowledge(
                    knowledge_name=knowledge_name,
                    description="错误代码数据知识库 - Maintained by 小胡"
                )
                if success:
                    logger.info(f"✅ 新知识库创建成功: {knowledge_name}")
                else:
                    logger.error(f"❌ 创建知识库失败: {knowledge_name}")
                    continue
            else:
                logger.info(f"📚 使用现有错误代码知识库: {knowledge_name} ({env} 环境)")
                dify_client.current_knowledge = dify_client.knowledge_list[knowledge_name]
            
            knowledge_id = dify_client.current_knowledge['id']
            logger.info(f"📋 错误代码知识库ID: {knowledge_id} ({env} 环境)")
            logger.info(f"📋 当前知识库信息: {dify_client.current_knowledge.get('name', 'Unknown')}")
            logger.info(f"🔗 使用的API端点: {dify_client.base_url}")
            
            # 上传错误代码文件
            for file_info in file_list:
                try:
                    results['total'] += 1
                    file_path = Path(file_info['file_path'])
                    document_name = file_info['file_name']
                    existing_doc_id = check_document_exists(dify_client, knowledge_id, document_name)
                    
                    if existing_doc_id:
                        dify_client.update_knowledge(
                            knowledge_id=knowledge_id,
                            document_id=existing_doc_id,
                            file_path=str(file_path)
                        )
                        action = "🔄 已更新"
                    else:
                        dify_client.upload_file(
                            knowledge_id=knowledge_id,
                            file_path=str(file_path)
                        )
                        action = "⬆️ 已上传"
                    
                    results['successful'] += 1
                    upload_type = "全量" if export_all else "增量"
                    logger.info(f"{action}错误代码({upload_type}): {document_name} ({env})")
                    
                except Exception as e:
                    error_msg = f"❌ 上传错误代码文件失败 {file_info['file_name']} ({env}): {str(e)}"
                    logger.error(error_msg)
                    results['errors'].append(error_msg)
                    
        except Exception as e:
            error_msg = f"❌ 处理 {env} 环境时出错: {str(e)}"
            logger.error(error_msg)
            logger.error(f"🔍 详细错误信息:\n{traceback.format_exc()}")
            results['errors'].append(error_msg)
    
    logger.info(f"🎉 错误代码数据上传完成！成功: {results['successful']}/{results['total']}")
    return results

def upload_to_faq_dify(data_dir=FAQ_DATA_DIR, environments=['local', 'aws'], export_all=True, files_to_upload=None):
    """
    根据产品名称上传FAQ数据到对应的Dify知识库
    现在支持两套知识库配置：
    1. FAQ-{product_name} - 使用OpenAI模型（原配置）
    2. BR-FAQ-{product_name} - 使用Amazon Bedrock Cohere模型（新配置）
    :param data_dir: FAQ数据目录
    :param environments: 要上传的环境列表
    :param export_all: 是否全量上传，False时进行增量上传
    :param files_to_upload: 增量上传时的文件列表
    :return: 上传结果统计
    """
    if export_all:
        logger.info("📋 开始全量上传FAQ数据到Dify（两套知识库配置）...")
    else:
        logger.info("🎯 开始增量上传FAQ数据到Dify（两套知识库配置）...")
    
    results = {'successful': 0, 'total': 0, 'errors': [], 'products': {}}
    
    # 根据export_all参数决定上传文件列表
    if export_all:
        # 全量上传：扫描目录获取所有文件
        data_path = Path(data_dir)
        if not data_path.exists():
            logger.error(f"❌ FAQ数据目录不存在: {data_path}")
            return results
            
        txt_files = list(data_path.glob("*.txt"))
        if not txt_files:
            logger.warning(f"⚠️ FAQ数据目录中没有找到txt文件: {data_path}")
            return results
            
        logger.info(f"📋 找到 {len(txt_files)} 个FAQ文件（全量上传）")
        file_list = txt_files
    else:
        # 增量上传：使用提供的文件列表
        if not files_to_upload:
            logger.info("😌 没有FAQ文件需要增量上传")
            return results
            
        # 过滤出FAQ文件
        faq_files = [f for f in files_to_upload if 'faq' in f['file_path'].lower()]
        if not faq_files:
            logger.info("😌 没有FAQ文件需要增量上传")
            return results
            
        logger.info(f"🎯 找到 {len(faq_files)} 个FAQ文件需要增量上传")
        # 转换为Path对象
        file_list = [Path(f['file_path']) for f in faq_files]
    
    # 按产品分组文件
    products = {}
    for txt_file in file_list:
        product_name = extract_product_name(txt_file.name)
        if product_name:
            if product_name not in products:
                products[product_name] = []
            products[product_name].append(txt_file)
        else:
            logger.warning(f"⚠️ 跳过无法识别产品名称的文件: {txt_file.name}")
    
    upload_type = "全量" if export_all else "增量"
    logger.info(f"📦 识别到 {len(products)} 个产品({upload_type}): {list(products.keys())}")
    results['products'] = {k: len(v) for k, v in products.items()}
    
    # 定义两套知识库配置
    knowledge_configs = [
        {
            'name_template': 'FAQ-{product_name}',
            'description_template': '{product_name} 产品FAQ和故障排除数据 - Maintained by 小胡',
            'handle_file_params': {
                'embedding_model': 'text-embedding-3-small',
                'embedding_model_provider': 'openai',
                'reranking_provider_name': 'cohere',
                'reranking_model_name': 'rerank-v3.5'
            }
        },
        {
            'name_template': 'BR-FAQ-{product_name}',
            'description_template': '{product_name} 产品FAQ和故障排除数据（Bedrock版本）- Maintained by 小胡',
            'handle_file_params': {
                'embedding_model': 'amazon.titan-embed-text-v2:0',
                'embedding_model_provider': 'bedrock',
                'reranking_provider_name': 'bedrock',
                'reranking_model_name': 'amazon.rerank-v1:0'
            }
        }
    ]
    
    # 遍历每个环境
    for env in environments:
        try:
            logger.info(f"🌍 开始处理 {env} 环境...")
            
            # 为每个产品和每个配置处理文件
            for product_name, files in products.items():
                for config in knowledge_configs:
                    try:
                        knowledge_name = config['name_template'].format(product_name=product_name)
                        description = config['description_template'].format(product_name=product_name)
                        handle_file_params = config['handle_file_params']
                        
                        logger.info(f"📚 处理产品知识库: {knowledge_name}")
                        
                        # 创建或获取产品知识库 - 避免使用类方法，使用独立客户端
                        product_client = FileToDify(url_data=DIFY_KNOWLEDGE_DATA[env])
                        product_client.query_knowledge_list()
                        
                        if knowledge_name not in product_client.knowledge_list:
                            logger.info(f"📚 创建新的产品知识库: {knowledge_name} ({env} 环境)")
                            # 使用实例方法创建知识库，避免类变量污染
                            success = product_client.create_knowledge(
                                knowledge_name=knowledge_name,
                                description=description
                            )
                            if not success:
                                logger.error(f"❌ 创建知识库失败: {knowledge_name}")
                                continue
                        else:
                            logger.info(f"📚 使用现有产品知识库: {knowledge_name} ({env} 环境)")
                            product_client.current_knowledge = product_client.knowledge_list[knowledge_name]
                        
                        knowledge_id = product_client.current_knowledge['id']
                        logger.info(f"📋 产品知识库ID: {knowledge_id} ({env} 环境), 当前知识库信息: {product_client.current_knowledge.get('name', 'Unknown')}")
                        
                        # 上传该产品的所有文件到当前知识库配置
                        for txt_file in files:
                            try:
                                results['total'] += 1
                                logger.info(f"📤 上传产品文件: {txt_file.name} 到 {knowledge_name}")
                                
                                # 检查文档是否已存在
                                document_name = txt_file.name
                                existing_doc_id = check_document_exists(product_client, knowledge_id, document_name)
                                
                                if existing_doc_id:
                                    # 更新现有文档
                                    product_client.update_knowledge(
                                        knowledge_id=knowledge_id,
                                        document_id=existing_doc_id,
                                        file_path=str(txt_file),
                                        **handle_file_params
                                    )
                                else:
                                    # 上传新文档
                                    product_client.upload_file(
                                        knowledge_id=knowledge_id,
                                        file_path=str(txt_file),
                                        **handle_file_params
                                    )
                                
                                results['successful'] += 1
                                upload_type = "全量" if export_all else "增量"
                                logger.info(f"✅ 产品文件上传成功({upload_type}): {txt_file.name} 到 {knowledge_name} ({env})")
                                
                            except Exception as e:
                                error_msg = f"❌ 上传产品文件失败 {txt_file.name} 到 {knowledge_name} ({env}): {str(e)}"
                                logger.error(error_msg)
                                results['errors'].append(error_msg)
                                
                    except Exception as e:
                        error_msg = f"❌ 处理产品 {product_name} 的知识库 {config['name_template']} 时出错 ({env}): {str(e)}"
                        logger.error(error_msg)
                        results['errors'].append(error_msg)
                    
        except Exception as e:
            error_msg = f"❌ 处理 {env} 环境时出错: {str(e)}"
            logger.error(error_msg)
            logger.error(f"🔍 详细错误信息:\n{traceback.format_exc()}")
            results['errors'].append(error_msg)
    
    logger.info(f"🎉 FAQ数据上传完成！成功: {results['successful']}/{results['total']}")
    return results

def upload_all_data(environments=['local', 'aws']):
    """
    上传所有数据到Dify（错误代码 + FAQ）
    :param environments: 要上传的环境列表
    :return: 合并的上传结果
    """
    logger.info("🚀 开始上传所有数据到Dify...")
    
    # 上传错误代码数据
    logger.info("=" * 50)
    errorcode_results = upload_to_errorcode_dify(environments=environments)
    
    # 上传FAQ数据
    logger.info("=" * 50)
    faq_results = upload_to_faq_dify(environments=environments)
    
    # 合并结果
    total_results = {
        'errorcode': errorcode_results,
        'faq': faq_results,
        'summary': {
            'total_successful': errorcode_results['successful'] + faq_results['successful'],
            'total_files': errorcode_results['total'] + faq_results['total'],
            'total_errors': len(errorcode_results['errors']) + len(faq_results['errors'])
        }
    }
    
    logger.info("=" * 50)
    logger.info("🎉 全部数据上传完成！")
    logger.info(f"📊 总体统计:")
    logger.info(f"  ✅ 成功上传: {total_results['summary']['total_successful']}")
    logger.info(f"  📁 总文件数: {total_results['summary']['total_files']}")
    logger.info(f"  ❌ 错误数量: {total_results['summary']['total_errors']}")
    logger.info(f"  🔧 错误代码文件: {errorcode_results['successful']}/{errorcode_results['total']}")
    logger.info(f"  📋 FAQ文件: {faq_results['successful']}/{faq_results['total']}")
    if faq_results['products']:
        logger.info(f"  📦 处理产品: {faq_results['products']}")
    
    return total_results

# 兼容性函数
def upload_all_txt_files():
    """向后兼容的函数"""
    results = upload_all_data()
    return results['summary']['total_successful'], results['summary']['total_files']



if __name__ == "__main__":
    logger.info("📤 直接运行Dify上传脚本...")
    try:
        results = upload_all_data()
        success_count = results['summary']['total_successful']
        total_count = results['summary']['total_files']
        logger.info(f"✅ Dify上传完成，成功上传 {success_count}/{total_count} 个文件")
    except Exception as e:
        logger.error(f"❌ Dify上传脚本运行失败: {str(e)}")
        logger.error(f"🔍 详细错误信息:\n{traceback.format_exc()}")